# https://gitee.com/luo_zhi_cheng/Mask2Former/blob/lzc/lzc/mask2former_triton/tools/trt_py_infer/trt_infer_fixed_exam.py

import cv2
import numpy as np
import pycuda.autoinit  # noqa: F401
import pycuda.driver as cuda
import tensorrt as trt
import time
import ctypes
from typing import List, Tuple


class YOLOV7_TRT(object):
    """
    description: A class that warps TensorRT ops.
    """

    def __init__(self, engine_file_path, plugin_path='', log_level=trt.Logger.INFO):
        print(trt.__version__)
        self.is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10  # is TensorRT >= 10
        # Create a Context on this device,
        self.ctx = cuda.Device(0).make_context()
        stream = cuda.Stream()
        TRT_LOGGER = trt.Logger(log_level) # INTERNAL_ERROR # ERROR # WARNING # INFO # VERBOSE
        runtime = trt.Runtime(TRT_LOGGER)

        # 加载自定义插件库
        # ctypes.CDLL(plugin_path)  # 使用ctypes强制加载
        # 初始化TensorRT插件系统
        trt.init_libnvinfer_plugins(TRT_LOGGER, "")

        # Deserialize the engine from file
        with open(engine_file_path, "rb") as f:
            engine = runtime.deserialize_cuda_engine(f.read())
        self.context = engine.create_execution_context()

        self.stream = stream
        # 去掉下面1行的 self.context = context ，在上面create_execution_context时候赋值给self.context，否则程序退出时就会出现下面类似的错误
        # [05/20/2025-01:36:08] [TRT] [E] 1: [defaultAllocator.cpp::deallocate::61] Error Code 1: Cuda Runtime (invalid argument)
        # [05/20/2025-01:36:08] [TRT] [E] 1: [cudaResources.cpp::~ScopedCudaStream::47] Error Code 1: Cuda Runtime (invalid device context)
        # [05/20/2025-01:36:08] [TRT] [E] 1: [cudaResources.cpp::~ScopedCudaEvent::24] Error Code 1: Cuda Runtime (context is destroyed)
        # self.context = context 

        # 提前分配内存等
        host_inputs = []
        cuda_inputs = []
        host_outputs = []
        cuda_outputs = []
        #trt8
        bindings = []
        #trt10
        input_binding_names = []
        output_binding_names = []
        input_shapes = []
        output_shapes = []

        for binding in engine:
            # binding 就是导出onnx时设置的网络的输入/输出的名称
            if self.is_trt10:
                binding_shape = engine.get_tensor_shape(binding)
                dtype = trt.nptype(engine.get_tensor_dtype(binding))
                is_input = (engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT)
            else:
                binding_shape = engine.get_binding_shape(binding)
                dtype = trt.nptype(engine.get_binding_dtype(binding))
                is_input = engine.binding_is_input(binding)
            print('binding_shape:', binding, binding_shape)
            self.batch_size = binding_shape[0]
            # trt.volume 是计算binding_shape各维度之乘积，但是binding_shape此时是动态维度，hw值为-1，结果与预期不符
            size = trt.volume(binding_shape)
            # Allocate host and device buffers
            # pagelocked_empty 用于开辟锁业内存，mem_alloc用于开辟显存
            # pagelocked_empty
            # 参数
            #     shape: 数组形状（如 (1024, 1024)）。
            #     dtype: 数据类型（如 numpy.float32）。
            #     order: 内存布局（"C" 为行优先，"F" 为列优先）。
            #     mem_flags: 可选标志（如 pycuda.driver.host_alloc_flags.DEVICEMAP 支持 GPU 直接访问）。
            # 返回
            #     numpy.ndarray对象
            host_mem = cuda.pagelocked_empty(size, dtype)
            # 参数
            #     bytes: 要分配的字节数（通常通过 nbytes 计算）
            # 返回
            #     pycuda._driver.DeviceAllocation对象，
            #     须知该对象通过 int() 转换后，会返回该对象对应的 ​​设备内存地址（GPU显存的物理地址）​​，以整数形式表示。
            #     这个地址是 CUDA 内核（kernel）或内存操作函数（如 memcpy）直接操作显存的关键标识
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            # Append the device buffer to device bindings.
            bindings.append(int(cuda_mem))
            # Append to the appropriate list.
            if is_input:
                input_binding_names.append(binding)
                input_shapes.append(binding_shape)
                host_inputs.append(host_mem)
                cuda_inputs.append(cuda_mem)
            else:
                output_binding_names.append(binding)
                output_shapes.append(binding_shape)
                host_outputs.append(host_mem)
                cuda_outputs.append(cuda_mem)

        # Store
        self.host_inputs = host_inputs
        self.cuda_inputs = cuda_inputs
        self.host_outputs = host_outputs
        self.cuda_outputs = cuda_outputs
        self.bindings = bindings
        self.input_binding_names = input_binding_names
        self.output_binding_names = output_binding_names
        self.input_shapes = input_shapes
        self.output_shapes = output_shapes

    def infer(self, net_inputs : np.ndarray):
        # Make self the active context, pushing it on top of the context stack.
        self.ctx.push()
        # Restore
        stream = self.stream
        context = self.context
        host_inputs = self.host_inputs
        cuda_inputs = self.cuda_inputs
        host_outputs = self.host_outputs
        cuda_outputs = self.cuda_outputs
        bindings = self.bindings
        input_binding_names = self.input_binding_names
        output_binding_names = self.output_binding_names
        input_shapes = self.input_shapes
        output_shapes = self.output_shapes

        # Copy input image to host buffer
        net_inputs = np.ascontiguousarray(net_inputs)
        np.copyto(host_inputs[0], net_inputs.ravel())

        # 推理infer
        # Transfer input data  to the GPU.
        for i in range(len(host_inputs)):
            cuda.memcpy_htod_async(cuda_inputs[i], host_inputs[i], stream)
        # Run inference.
        if self.is_trt10:
            for i in range(len(cuda_inputs)):
                context.set_tensor_address(input_binding_names[i], cuda_inputs[i])
            for i in range(len(cuda_outputs)):
                context.set_tensor_address(output_binding_names[i], cuda_outputs[i])
            context.execute_async_v3(stream_handle=stream.handle)
        else:
            context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
        # Transfer predictions back from the GPU.
        for i in range(len(host_outputs)):
            cuda.memcpy_dtoh_async(host_outputs[i], cuda_outputs[i], stream)
        # Synchronize the stream
        stream.synchronize()
        # Remove any context from the top of the context stack, deactivating it.
        self.ctx.pop()

        output_rsp = []
        for i in range(len(host_outputs)):
            output_rsp.append(host_outputs[i].reshape(output_shapes[i]))

        return output_rsp
        
    def destroy(self):
        # Remove any context from the top of the context stack, deactivating it.
        self.ctx.pop()


def infer_trt():
    engine_file_path = "/weights/chuwei_best_20240805_fp16_b1.trt"
    # PLUGIN_LIBRARY = "xxx.so"  # 替换为你的插件库路径
    model_wrapper = YOLOV7_TRT(engine_file_path)

    try:
        import sys
        import os.path as osp
        sys.path.append(osp.dirname(osp.dirname(osp.dirname(osp.abspath(__file__)))))
        from lzc.new_files.preprocess import preprocess
        img_path = '/data/chuwei/toilet_30.jpg'
        img = cv2.imread(img_path)
        img_pre, ratio, dwdh = preprocess(img)
        net_outputs = model_wrapper.infer(img_pre)
    finally:
        # destroy the instance
        model_wrapper.destroy()

    return net_outputs


if __name__ == "__main__":
    pred = infer_trt()
    print('finish')
